# utils/system_check.py
import torch
import psutil
import platform
from pathlib import Path

class SystemChecker:
    """系统检查工具"""
    
    def __init__(self):
        self.checks = []
    
    def check_pytorch(self):
        """检查PyTorch"""
        try:
            cuda_available = torch.cuda.is_available()
            cuda_version = torch.version.cuda if cuda_available else "N/A"
            gpu_count = torch.cuda.device_count() if cuda_available else 0
            
            self.checks.append({
                "component": "PyTorch",
                "status": "OK",
                "details": f"版本: {torch.__version__}, CUDA: {cuda_available}, GPU数量: {gpu_count}"
            })
            
            return True
        except Exception as e:
            self.checks.append({
                "component": "PyTorch", 
                "status": "ERROR",
                "details": str(e)
            })
            return False
    
    def check_memory(self):
        """检查内存"""
        try:
            memory = psutil.virtual_memory()
            disk = psutil.disk_usage('/')
            
            self.checks.append({
                "component": "内存",
                "status": "OK", 
                "details": f"内存: {memory.percent}% 使用, 磁盘: {disk.percent}% 使用"
            })
            
            return memory.percent < 90 and disk.percent < 90
        except Exception as e:
            self.checks.append({
                "component": "内存",
                "status": "ERROR",
                "details": str(e)
            })
            return False
    
    def check_directories(self):
        """检查目录结构"""
        required_dirs = [
            "data/training",
            "data/models", 
            "data/game_state",
            "logs",
            "visualizations"
        ]
        
        missing_dirs = []
        for dir_path in required_dirs:
            if not Path(dir_path).exists():
                missing_dirs.append(dir_path)
        
        if missing_dirs:
            self.checks.append({
                "component": "目录结构",
                "status": "WARNING",
                "details": f"缺少目录: {', '.join(missing_dirs)}"
            })
            return False
        else:
            self.checks.append({
                "component": "目录结构",
                "status": "OK",
                "details": "所有必需目录都存在"
            })
            return True
    
    def check_system(self):
        """执行系统检查"""
        print("执行系统检查...")
        
        checks_passed = all([
            self.check_pytorch(),
            self.check_memory(), 
            self.check_directories()
        ])
        
        # 打印检查结果
        for check in self.checks:
            status_icon = "✅" if check["status"] == "OK" else "⚠️" if check["status"] == "WARNING" else "❌"
            print(f"{status_icon} {check['component']}: {check['details']}")
        
        return checks_passed
    
    def full_system_check(self):
        """完整系统检查"""
        print("=" * 50)
        print("翁法罗斯平台 - 完整系统检查")
        print("=" * 50)
        
        # 系统信息
        print(f"操作系统: {platform.system()} {platform.release()}")
        print(f"Python版本: {platform.python_version()}")
        print(f"处理器: {platform.processor()}")
        
        # 执行检查
        self.check_system()
        
        # 建议
        print("\n建议:")
        if not torch.cuda.is_available():
            print("⚠️ 未检测到CUDA，训练速度可能较慢")
        
        memory = psutil.virtual_memory()
        if memory.percent > 80:
            print("⚠️ 内存使用率较高，建议关闭其他应用")